{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "32ed78ad-e5f1-46f7-80e5-6657d736b263",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.7.0+cu126\n",
      "12.6\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "#cuda版本的代码是在家里电脑上运行的结果\n",
    "import torch\n",
    "print(torch.__version__)\n",
    "print(torch.version.cuda)\n",
    "print(torch.cuda.is_available())  #输出为True，则安装无误\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e683b53c-7129-4c1f-a9a4-17778c3e11f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'torch.Tensor'> tensor([ 1, 12, 20])\n",
      "<class 'torch.Tensor'> tensor(2)\n",
      "<class 'torch.Tensor'> tensor([ 3, 14, 22])\n",
      "<class 'torch.Tensor'> tensor([ 2, 24, 40])\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor([1,12,20])\n",
    "print(type(a),a)\n",
    "\n",
    "b = torch.tensor(2)\n",
    "print(type(a),b)\n",
    "\n",
    "c = a+b\n",
    "print(type(c),c)\n",
    "\n",
    "d = a*b\n",
    "print(type(d),d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b75aee66-e66e-4dd8-813f-dda39c6520c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "torch.device('cpu'),torch.device('cuda'),torch.device('cuda:1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b82b829c-6712-4af5-b994-745f6e758d20",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.device_count()  #查询GPU核心的数量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b17d48b7-abf0-46ab-9819-5f67b1abe8f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "def try_gpu(i=0):\n",
    "    #如果存在gpu则返回gpu[i],否则返回cpu\n",
    "    if torch.cuda.device_count() >= i+1:\n",
    "        return torch.device(f'cuda{i}')\n",
    "    return torch.device('cpu')         \n",
    "    \n",
    "    \n",
    "x = torch.tensor([5.67],device=try_gpu())\n",
    "print(x.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4b39a9ff-1876-435a-94d0-457ead146b69",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'torch' from 'D:\\\\ProgramData\\\\anaconda3\\\\Lib\\\\site-packages\\\\torch\\\\__init__.py'>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "net = torch.nn.Sequential(torch.nn.Linear(3,1))\n",
    "net = net.to(device=try_gpu())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06e603b8-495d-48b0-b675-0ddc67f82b55",
   "metadata": {},
   "outputs": [],
   "source": [
    "net[0].weight.data.device"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
